use std::io::Result;
use std::sync::mpsc::channel;

use image::{ImageBuffer, Rgb};
use num::{Complex, ToPrimitive};
use threadpool::ThreadPool;

fn julia(c: Complex<f64>, x: u32, y: u32, width: u32, height: u32, iterations: u32) -> f64 {
    let mut z = Complex::new(
        x as f64 / width as f64 * 4.0 - 2.0,
        y as f64 / height as f64 * 4.0 - 2.0,
    );
    let mut n = 0;
    while n < iterations && z.re * z.re + z.im * z.im < 4.0 {
        z = z.sqrt() + c;
        n += 1;
    }
    n.to_f64().unwrap()
}

fn wavelength_to_rgb(wavelength: f64) -> Rgb<u8> {
    let mut r = 0.0;
    let mut g = 0.0;
    let mut b = 0.0;

    if wavelength >= 380.0 && wavelength <= 440.0 {
        r = -(wavelength - 440.0) / (440.0 - 380.0);
        g = 0.0;
        b = 1.0;
    } else if wavelength >= 440.0 && wavelength <= 490.0 {
        r = 0.0;
        g = (wavelength - 440.0) / (490.0 - 440.0);
        b = 1.0;
    } else if wavelength >= 490.0 && wavelength <= 510.0 {
        r = 0.0;
        g = 1.0;
        b = -(wavelength - 510.0) / (510.0 - 490.0);
    } else if wavelength >= 510.0 && wavelength <= 580.0 {
        r = (wavelength - 510.0) / (580.0 - 510.0);
        g = 1.0;
        b = 0.0;
    } else if wavelength >= 580.0 && wavelength <= 645.0 {
        r = 1.0;
        g = -(wavelength - 645.0) / (645.0 - 580.0);
        b = 0.0;
    } else if wavelength >= 645.0 && wavelength <= 780.0 {
        r = 1.0;
        g = 0.0;
        b = 0.0;
    }

    // 线性到sRGB转换
    let r = if r > 0.04045 {
        ((r + 0.055) / 1.055).powf(2.4)
    } else {
        r / 12.92
    };
    let g = if g > 0.04045 {
        ((g + 0.055) / 1.055).powf(2.4)
    } else {
        g / 12.92
    };
    let b = if b > 0.04045 {
        ((b + 0.055) / 1.055).powf(2.4)
    } else {
        b / 12.92
    };

    // 转换为u8
    let r = (r * 255.0) as u8;
    let g = (g * 255.0) as u8;
    let b = (b * 255.0) as u8;

    Rgb([r, g, b])
}

fn main() -> Result<()> {
    let (width, height) = (1920, 1080);
    let mut img = ImageBuffer::new(width, height);
    let iterations = 300;

    let c = Complex::new(-0.8, 0.156);

    let pool = ThreadPool::new(num_cpus::get());
    let (tx, rx) = channel();

    for y in 0..height {
        let tx = tx.clone();
        pool.execute(move || {
            for x in 0..width {
                let i = julia(c, x, y, width, height, iterations);
                let pixel = wavelength_to_rgb(380.0 + i * 400.0 / iterations.to_f64().unwrap());
                tx.send((x, y, pixel)).expect("Could not send data!");
            }
        });
    }

    for _ in 0..(width * height) {
        let (x, y, pixel) = rx.recv().expect("can't compute");
        img.put_pixel(x, y, pixel);
    }

    let _ = img
        .save("resource/thread/draw/output.png")
        .expect("can't save");

    Ok(())
}
